from datetime import datetime
import backtrader as bt
import matplotlib.pyplot as plt
import akshare as ak
import pandas as pd

plt.rcParams["font.sans-serif"] = ["SimHei"]
plt.rcParams["axes.unicode_minus"] = False

# 获取股票后复权数据
stock_hfq_df = ak.stock_zh_a_hist(symbol="000001", adjust="hfq").iloc[:, :6]
# 处理字段命名
stock_hfq_df.columns = ['date', 'open', 'close', 'high', 'low', 'volume']
# 将date设为索引
stock_hfq_df.index = pd.to_datetime(stock_hfq_df['date'])

class MyStrategy(bt.Strategy):
    """
    主策略程序
    """
    params = (
        ("fast_period", 5),  # 短期（5日）移动平均线周期
        ("slow_period", 30),  # 长期（30日）移动平均线周期
    )

    def __init__(self):
        self.data_close = self.datas[0].close
        self.order = None
        self.buy_price = None
        self.buy_comm = None
        
        # 初始化短期和长期移动均线指标
        self.fast_sma = bt.indicators.SimpleMovingAverage(self.datas[0], period=self.params.fast_period)
        self.slow_sma = bt.indicators.SimpleMovingAverage(self.datas[0], period=self.params.slow_period)

    def next(self):
        if self.order:  
            return
        
        # 检查金叉和死叉
        if self.fast_sma[0] > self.slow_sma[0] and self.fast_sma[-1] <= self.slow_sma[-1]:  # 金叉，买入
            #self.order = self.buy(size=100)
            self.order_target_percent(target=0.10)
        elif self.fast_sma[0] < self.slow_sma[0] and self.fast_sma[-1] >= self.slow_sma[-1]:  # 死叉，卖出
            #self.order = self.sell(size=100)
            self.order_target_percent(target=0.10)

# 初始化回测系统
cerebro = bt.Cerebro()

# 设置回测时间范围
start_date = datetime(2010, 4, 3)
end_date = datetime(2024, 5, 22)

# 加载数据到Cerebro
data = bt.feeds.PandasData(dataname=stock_hfq_df, fromdate=start_date, todate=end_date)
cerebro.adddata(data)

# 加载策略到Cerebro
cerebro.addstrategy(MyStrategy)

# 设置初始资金和手续费
start_cash = 1000000
cerebro.broker.setcash(start_cash)
cerebro.broker.setcommission(commission=0.002)

# 运行回测
cerebro.run()

# 获取回测结果并打印
port_value = cerebro.broker.getvalue()
pnl = port_value - start_cash
print(f"初始资金: {start_cash}\n回测期间：{start_date.strftime('%Y%m%d')}:{end_date.strftime('%Y%m%d')}")
print(f"总资金: {round(port_value, 2)}")
print(f"净收益: {round(pnl, 2)}")

# 绘制图表
cerebro.plot(style='candlestick',figsize=(12, 8))
